import psycopg2
import io
from ..base.module import BaseANN
import numpy as np


def get_conn():
    conn = psycopg2.connect(host="127.0.0.1", user="ann", password="Huawei@123", dbname="ann", port=6437)
    conn.autocommit = True
    return conn

class OpengaussHuaweiVector(BaseANN):
    def __init__(self, metric, m, efConstruction):
        super().__init__()
        self._m = m
        self._efConstruction = efConstruction
        self._metric = metric
        self._cur = None
        self.parallel_workers = 8  # 默认值
        self.conn = None
        self.conn = get_conn()
        self._cur = self.conn.cursor()
        self._ef_search = None

    
    def _update_query(self):
        """根据当前的metric和tablename更新查询语句"""
        if self._metric == "angular":
            self._query = f"SELECT id FROM {self.tablename} ORDER BY embedding <=> %s LIMIT %s"
        elif self._metric == "euclidean":
            self._query = f"SELECT id FROM {self.tablename} ORDER BY embedding <-> %s LIMIT %s"
        else:
            raise RuntimeError(f"unknown metric {self._metric}")
    
    def get_row_count(self) -> int:
        """获取表中的行数
        重写BaseANN中的抽象方法，返回表中的行数
        """
        self._cur.execute(f"SELECT COUNT(*) FROM {self.tablename}")
        return self._cur.fetchone()[0]

    def get_vector_by_id(self, vector_id):
        """根据ID获取向量
        重写BaseANN中的抽象方法，从数据库获取指定ID的向量
        
        Args:
            vector_id: 向量的ID
            
        Returns:
            numpy.ndarray: 向量数据，如果找不到则返回None
        """
        try:
            # 从数据库获取对应的向量
            self._cur.execute(f"SELECT embedding::text FROM {self.tablename} WHERE id = {vector_id}")
            db_result = self._cur.fetchone()
            if db_result is None:
                return None
            
            db_vector = db_result[0]
            
            # 清理字符串并提取数值
            db_str = db_vector.strip()
            # 移除括号，支持[]和()格式
            if (db_str.startswith('[') and db_str.endswith(']')) or \
               (db_str.startswith('(') and db_str.endswith(')')):
                db_str = db_str[1:-1]
            
            # 分割并转换为浮点数
            return np.array([float(x.strip()) for x in db_str.split(',') if x.strip()])
        except Exception as e:
            print(f"获取向量出错: {e}")
            return None

    def has_same_data(self, dataset) -> bool:
        """重用父类BaseANN中的通用实现"""
        return super().has_same_data(dataset)

    def done(self):
        self.drop()
        self.conn.close()

    def drop(self):
        """Drop existing table"""
        self._cur.execute(f"DROP TABLE IF EXISTS {self.tablename};")

    def copy(self, dataset):
        """Copy data to table using batch processing
        Args:
            dataset: h5py Dataset object
        """
        # 获取训练数据集
        train_data = dataset["train"]

        print("copying data...")
        # 创建表
        self._cur.execute(f"CREATE TABLE {self.tablename} (id int, embedding vector(({train_data.shape[1]})))")
        self._cur.execute(f"ALTER TABLE {self.tablename} ALTER COLUMN embedding SET STORAGE PLAIN")

        # 批量处理
        batch_size = 10000
        total_rows = train_data.shape[0]

        print(f"Start copying {total_rows} rows with batch size {batch_size}...")
        for start_idx in range(0, total_rows, batch_size):
            end_idx = min(start_idx + batch_size, total_rows)
            batch = train_data[start_idx:end_idx]  # 只读取一部分数据

            # 使用 StringIO 进行批量写入
            data = io.StringIO()
            for i, embedding in enumerate(batch):
                embedding_str = "[" + ", ".join(f"{x:.8f}".rstrip("0").rstrip(".") for x in embedding) + "]"
                data.write(f"{start_idx + i}\t{embedding_str}\n")

            data.seek(0)
            self._cur.copy_from(data, self.tablename, columns=('id', 'embedding'), sep='\t')

            if end_idx % (total_rows // 10) == 0 or end_idx == total_rows:
                print(f"Processed {end_idx}/{total_rows} rows ({(end_idx / total_rows * 100):.2f}%)")

    def has_same_index(self) -> bool:
        """检查是否有相同参数的索引"""
        try:
            # 1. 检查表是否存在
            self._cur.execute(f"""
                SELECT COUNT(*) 
                FROM pg_tables 
                WHERE tablename = '{self.tablename}'
            """)
            if self._cur.fetchone()[0] == 0:
                print(f"表 {self.tablename} 不存在")
                return False


            # 2. 检查当前索引是否存在
            self._cur.execute(f"""
                SELECT COUNT(*) 
                FROM pg_indexes 
                WHERE tablename = '{self.tablename}' 
                AND indexname = '{self.index_name}'
            """)
            if self._cur.fetchone()[0] == 0:
                print(f"索引 {self.index_name} 不存在")
                return False
            
            # 3. 删除其它索引
            self.drop_other_index()
                
            print(f"索引 {self.index_name} 校验通过，可以重用现有索引")
            return True

        except Exception as e:
            print(f"检查索引时出错: {e}")
            return False

    def drop_other_index(self):
        self._cur.execute(f"""
        SELECT indexname
        FROM pg_indexes
        WHERE tablename = '{self.tablename}'
        AND indexname != '{self.index_name}'
        """)
        conflicting_indexes = self._cur.fetchall()
        if conflicting_indexes:
            index_names = [row[0] for row in conflicting_indexes]
            print(f"表 {self.tablename} 上存在其他索引，准备删除: {', '.join(index_names)}")

            for row in conflicting_indexes:
                index_name = row[0]
                print(f"删除索引 {index_name}")
                self._cur.execute(f"DROP INDEX IF EXISTS {index_name}")


    def create_index(self):
        """Create index"""
        # 创建索引前，先删除索引
        self.drop_other_index()
        self._cur.execute(f"DROP INDEX IF EXISTS {self.index_name}")
        self._cur.execute(f"ALTER TABLE items SET(parallel_workers={min(self.parallel_workers, 32)});")
        print(f"创建索引 {self.index_name}...")
        if self._metric == "angular":
            self._cur.execute(
                f"CREATE INDEX {self.index_name} ON {self.tablename} USING hnsw (embedding vector_cosine_ops) WITH (m = {self._m}, ef_construction = {self._efConstruction})"
            )
        elif self._metric == "euclidean":
            self._cur.execute(
                f"CREATE INDEX {self.index_name} ON {self.tablename} USING hnsw (embedding vector_l2_ops) WITH (m = {self._m}, ef_construction = {self._efConstruction})"
            )
        else:
            raise RuntimeError(f"unknown metric {self._metric}")
        print("done!")

    def set_query_arguments(self, ef_search):
        self._ef_search = ef_search
        self._cur.execute(f"SET enable_seqscan TO off; SET hnsw_ef_search = {self._ef_search}")

    def query(self, v, n):
        v_str = "[" + ", ".join(f"{x:.8f}" for x in v) + "]"
        self._cur.execute(self._query, (v_str, n))
        res = [id for id, in self._cur.fetchall()]
        return res

    def get_memory_usage(self):
        if self._cur is None:
            return 0
        self._cur.execute(f"SELECT pg_total_relation_size('{self.index_name}')")
        return self._cur.fetchone()[0] / 1024

    def get_table_usage(self):
        self._cur.execute(f"SELECT pg_table_size('{self.tablename}')")
        return self._cur.fetchone()[0] / 1024

    def __str__(self):
        return f"opengauss-huawei-hnsw(m={self._m}, ef_construction={self._efConstruction}, ef_search={self._ef_search})"

class OpenGaussHuaweiIvfflatVector(BaseANN):
    def __init__(self, metric, nlist):
        super().__init__()
        self._n_list = nlist
        self._metric = metric
        
    @property
    def index_name(self):
        """获取IVFFlat索引名称"""
        return f"{self.tablename}_ogh_ivfflat_{self._metric}_{self._n_list}_idx"

    def create_index(self):
        """Create IVFFlat index"""
        # 创建索引前，先删除索引
        self.drop_other_index()
        self._cur.execute(f"DROP INDEX IF EXISTS {self.index_name}")
        print(f"创建IVFFlat索引 {self.index_name}...")
        if self._metric == "angular":
            self._cur.execute(
                f"CREATE INDEX {self.index_name} ON {self.tablename} USING ivfflat (embedding vector_cosine_ops) WITH (lists = {self._n_list})"
            )
        elif self._metric == "euclidean":
            self._cur.execute(
                f"CREATE INDEX {self.index_name} ON {self.tablename} USING ivfflat (embedding vector_l2_ops) WITH (lists = {self._n_list})"
            )
        else:
            raise RuntimeError(f"unknown metric {self._metric}")
        print("done!")

class OpenGaussHuaweiHnswVector(BaseANN):
    def __init__(self, metric, m, efConstruction):
        super().__init__()
        self._metric = metric
        self._m = m
        self._ef_construction = efConstruction
        
    @property
    def index_name(self):
        """获取HNSW索引名称"""
        return f"{self.tablename}_ogh_hnsw_{self._metric}_{self._m}_{self._ef_construction}_idx"

    def create_index(self):
        """Create HNSW index"""
        # 创建索引前，先删除索引
        self.drop_other_index()
        self._cur.execute(f"DROP INDEX IF EXISTS {self.index_name}")
        print(f"创建HNSW索引 {self.index_name}...")
        if self._metric == "angular":
            self._cur.execute(
                f"CREATE INDEX {self.index_name} ON {self.tablename} USING hnsw (embedding vector_cosine_ops) WITH (m = {self._m}, ef_construction = {self._ef_construction})"
            )
        elif self._metric == "euclidean":
            self._cur.execute(
                f"CREATE INDEX {self.index_name} ON {self.tablename} USING hnsw (embedding vector_l2_ops) WITH (m = {self._m}, ef_construction = {self._ef_construction})"
            )
        else:
            raise RuntimeError(f"unknown metric {self._metric}")
        print("done!")


